# -*- coding: UTF-8 -*-
# Copyright (c) 2023 Baidu.com, Inc. All Rights Reserved
"""
本文件实现了单句询问模型，输出格式为：
姓名      问题1     问题2
虚拟人1    回复1     回复2

目前屏蔽了后处理；支持输出每句耗时，需要则解除注释"worksheet.append(output_time_list)"

Authors: yanghao31(ouhaojie@baidu.com)
Date:    2023/11/9 17:13
"""

import os
import requests
import xlrd2
import xlwt
import time
import json
import openpyxl as op
from openpyxl import load_workbook
import openpyxl
import re
from tqdm import tqdm
MY_DIR = os.path.abspath(".")

INPUT_FILE = "英语人设信息评测集.xlsx"
INFO_FILE = "人设信息(json格式)-v4.xlsx"
OUTPUT_FILE = MY_DIR + '/output_{}_{}.xlsx'.format(INPUT_FILE.split(".")[0],
                                                  time.strftime("%Y-%m-%d %H_%M_%S", time.localtime(time.time())))
# INPUT_FILE = "dynamic_evaluate_topic_pm.xlsx"
PRE_PROCESS_URL = "http://jp02-gbu-k8s-a30-033.jp02.baidu.com:8000/aichat/en/preprocess"
# 线下32L
MODEL_URL = "http://jp02-gbu-k8s-a30-033.jp02.baidu.com:8001/aichat/en/model"
# 线下32L
# MODEL_URL = "http://jp02-gbu-k8s-a30-033.jp02.baidu.com:8011/aichat/en/model"
# 线下llama2
MODEL_URL = "http://jp02-gbu-k8s-a30-033.jp02.baidu.com:8400/aichat/en/model"
# 预发布llama2
MODEL_URL = "http://gbu.jp02-a30-apisix-sandbox.baidu-int.com/inner-model/aichat/en/model"
# 线上llama2
MODEL_URL = "http://gbu.jp02-a30-apisix-online.baidu-int.com/inner-model/aichat/en/model"
# 英语后处理
POST_PROCESS_URL = "http://jp02-gbu-k8s-a30-033.jp02.baidu.com:8002/aichat/en/postprocess"

WHO = [
    "can you tell me your name"
]
AGE = [
    "what is your age",
]
SEX = [
    "tell me your gender"
]


def input_init(input_file):
    """
    第一个sheet从上到下，再从左到右
    :param input_file:
    :return:
    """
    input_list = []
    wb = load_workbook(input_file)
    sheet0 = wb[wb.sheetnames[0]]
    cols = sheet0.max_column
    for j in range(cols):
        rows = sheet0.max_row
        for i in range(rows):
            cell = sheet0.cell(row=i + 2, column=j + 1).value
            if cell is not None:
                input_list.append(cell)
    # print(input_list)
    return input_list


def init_output(output_file):
    # 判断文件是否存在
    if not os.path.exists(output_file):
        # 如果文件不存在，创建新的工作簿并选择第一个工作表
        workbook = openpyxl.Workbook()
        worksheet = workbook.active

        # 写入表头内容
        header = ['姓名']
        wb = load_workbook(INPUT_FILE)
        sheet0 = wb[wb.sheetnames[0]]
        rows = sheet0.max_row
        for i in range(rows):
            cell = sheet0.cell(row=i + 2, column=1).value
            if cell is not None:
                header.append(f"{i + 1}:{cell}")
        worksheet.append(header)

        # 保存工作簿
        workbook.save(output_file)


def en_pre_process(context, bot_name, previous_content_list, bot_gender=0):
    """
    调用英语模型
    :param bot_name:
    :param previous_content_list: 候选集
    :param context: 内容
    :param bot_gender: 虚拟人性别，当传递入参ext_info的时候，机器人性别代码：0表示男，1表示女；用户性别代码：1代表男，2代表女，3代表其他
    :return:
    """
    # 获取参数
    method = "POST"

    get_params = {}
    post_params = {
        "context": context,
        "user_id": "1000001004&1534976200",
        "ext_info": {
            "user_age": 21,
            "user_gender": 2 - bot_gender,
            "conversation_cnt": 3,
            "user_name": "john",
            "language": "en",
            "bot_name": bot_name,
            "bot_age": 27,
            "bot_gender": bot_gender,
            "module_call_cnt": {
                "en_model": 3,
                "en_postprocess": 2,
                "en_preprocess": 3
            },
            "previous_content": previous_content_list
        },
        "state": "start"
    }
    # print(post_params)

    url = PRE_PROCESS_URL

    headers = {"Content-Type": "application/json", "X-REQUEST-ID": "111",
               "Cookie": "BIDUPSID=C0ADD4B031859AB61D8EA66BF1BE81C7; PSTM=1653481521"}
    try:
        resp = requests.request(method, url, params=get_params, json=post_params, headers=headers)
        # print(resp.text)
        assert resp.status_code == 200
    except Exception as e:
        # return {'code': -1, 'msg': str(e), 'result': False}
        print(e)
        return -1
    return resp


def en_model(context, bot_name, previous_content_list, bot_gender=0):
    """
    调用英语模型
    :param bot_name:
    :param previous_content_list: 多消息过滤
    :param context: 内容
    :param bot_gender: 虚拟人性别，当传递入参ext_info的时候，机器人性别代码：0表示男，1表示女；用户性别代码：1代表男，2代表女，3代表其他
    :return:
    """
    # 获取参数
    method = "POST"

    get_params = {}
    post_params = {
        "user_id": "1531033666_1531090035",
        "context": context,
        "ext_info": {
            "user_age": 21,
            "user_gender": 2 - bot_gender,
            "plato_gender": 2,
            "conversation_cnt": 3,
            "user_name": "john",
            "language": "en",
            "bot_name": bot_name,
            "bot_age": 27,
            "bot_gender": bot_gender,
            "module_call_cnt": {
                "en_model": 3,
                "en_postprocess": 2,
                "en_preprocess": 3
            },
            "previous_content": previous_content_list
        }
    }
    # print(post_params)

    url = MODEL_URL

    headers = {"Content-Type": "application/json",
               "Cookie": "BAIDUID=8305EA057A02D18DB9B3886DBA4ADB9D:FG=1;"
                         " BIDUPSID=8305EA057A02D18DF6A09631D64D0455; PSTM=1654677176",
               "x-request-id": "123",
               "mode": "en_llm_service"
               }
    try:
        start_time = time.time()  # 获取当前时间戳
        resp = requests.request(method, url, params=get_params, json=post_params, headers=headers)
        end_time = time.time()  # 获取当前时间戳
        resp_time = end_time - start_time  # 计算接口返回耗时，单位为秒
        # print(resp.text)
        assert resp.status_code == 200
        # print(resp.text)
    except Exception as e:
        # return {'code': -1, 'msg': str(e), 'result': False}
        return -1
    return resp, resp_time


def en_post_process(context, bot_name, candidates_list, previous_content_list, bot_gender=0):
    """
    调用英语模型
    :param bot_name:
    :param candidates_list:
    :param previous_content_list: 候选集
    :param context: 内容
    :param bot_gender: 虚拟人性别，当传递入参ext_info的时候，机器人性别代码：0表示男，1表示女；用户性别代码：1代表男，2代表女，3代表其他
    :return:
    """
    # 获取参数
    method = "POST"

    get_params = {}
    post_params = {
        "context": context,
        "candidates": candidates_list,
        "user_id": "1000001004&1534976200",
        "ext_info": {
            "user_age": 21,
            "user_gender": 2 - bot_gender,
            "conversation_cnt": 3,
            "user_name": "john",
            "language": "en",
            "bot_name": bot_name,
            "bot_age": 27,
            "bot_gender": bot_gender,
            "module_call_cnt": {
                "en_model": 3,
                "en_postprocess": 2,
                "en_preprocess": 3
            },
            "previous_content": previous_content_list
        },
        "state": "start"
    }
    # print(post_params)

    url = POST_PROCESS_URL

    headers = {"Content-Type": "application/json", "X-REQUEST-ID": "111",
               "Cookie": "BIDUPSID=C0ADD4B031859AB61D8EA66BF1BE81C7; PSTM=1653481521"}
    try:
        resp = requests.request(method, url, params=get_params, json=post_params, headers=headers)
        # print(resp.text)
        assert resp.status_code == 200
    except Exception as e:
        # return {'code': -1, 'msg': str(e), 'result': False}
        return -1
    return resp


def work(input_list, name, info_prompt):
    global totalTime,totalQuesNum
    # 将结果写入Excel中
    workbook = openpyxl.load_workbook(OUTPUT_FILE)

    # 写入具体内容
    input_list = input_list
    # input_list = input_init(INPUT_FILE)
    worksheet = workbook.worksheets[0]
    output_tmp_list = [name]
    output_time_list = ["请求时长"]
    print('对虚拟人设{},进行人设稳定性测试提问'.format(name))
    for i in range(len(input_list)):
        age,name,gender,job = '','','',''
        question = [{
                "text": input_list[i],
                "role_type": 1,
                "source": "user",
                "lang": "en"
            }]
        context = info_prompt + question
        previous_content_list = []

        pre_result = en_pre_process(context=context, bot_name=name, previous_content_list=previous_content_list)
        if pre_result == -1:
            output_tmp_list.append(f"{i + 1}:pre_process error")
            output_time_list.append(f"{i + 1}:pre_process error")

            continue

        candidates_list = json.loads(pre_result.text).get("data").get("stage_info").get("candidates")
        # print(candidates_list)
        time.sleep(0.1)

        if len(candidates_list) == 0:
            startTime=time.time()
            result, result_time = en_model(context=context, bot_name=name, previous_content_list=previous_content_list)
            if result == -1:
                output_tmp_list.append(f"{i + 1}:model error")
                output_time_list.append(f"{i + 1}:model error")
                continue
            # print(result.text)
            try:
                candidates_list = json.loads(result.text).get("data").get("candidates")
                print("{}.{} candidates: {}".format(i, input_list[i], candidates_list[0].get("text")))
                userTime=time.time()-startTime
                totalTime+=userTime
                totalQuesNum+=1
            except Exception as e:
                print(e)
                print('Error:'+result.text)
                continue


            if i==0:
                text = candidates_list[0].get("text").replace(',', '.').replace('!', '.')
                name1 = re.findall("name is (.*?)\.", text)
                name3=re.findall("I'm (.*?)\.",text)
                name2 = re.findall("I am (.*?)\.", text)
                if len(name1):
                    name=name1[0]
                elif len(name2):
                    name = name2[0]
                elif len(name3):
                    name = name3[0]
                else:
                    name=''
            if i == 1:

                text = candidates_list[0].get("text").replace(',', '.')
                try:
                    age = re.findall("\d+", text)[0]
                except Exception as e:
                    age = ''

            if i==2:
                text = candidates_list[0].get("text").replace(',', '.').replace('!', '.').replace('an ', '').replace('a ', '')
                gender1=re.findall("I'm (.*?)\.",text)
                gender2 = re.findall("I am (.*?)\.", text)

                if len(gender1):
                    gender=gender1[0]
                elif len(gender2):
                    gender = gender2[0]
                else:
                    gender=''
            if i == 3:
                text = candidates_list[0].get("text").replace(',', '.').replace('!', '.').replace('an ', '').replace(
                    'a ', '')
                job1 = re.findall("I'm (.*?)\.", text)
                job2 = re.findall("I am (.*?)\.", text)

                if len(job1):
                    job = job1[0]
                elif len(job2):
                    job = job2[0]
                else:
                    job = ''
                jobSpilt=job.split(' ')
                job=' '.join(jobSpilt[:2])

        else:
            output_tmp_list.append(f"{i + 1}:触发前处理")
            output_time_list.append(f"{i + 1}:触发前处理")
        time.sleep(0.1)

        output_tmp_list.append(candidates_list[0].get("text"))

        output_time_list.append(result_time)
        # 后处理
        # post_result = en_post_process(context=context, bot_name=name, candidates_list=candidates_list,
        #                               previous_content_list=previous_content_list)
        # # print(result)
        # if post_result == -1:
        #     output_tmp_list.append(f"{i}:post_process error")
        # reply = "en_post_process error"
        # reply = json.loads(post_result.text).get("data").get("reply").get("text")
        # print("{}.{} reply: {}".format(i, name, reply))
        # output_tmp_list.append(reply)
        # if reply != candidates_list[0].get("text"):
        #     output_tmp_list.append(f"{i}:触发后处理")
        # previous_content_list.append(reply)

    output_tmp_list.append(json.dumps({'name':name,'age':age,'gender':gender,'job':job}))


    print('提取llama回复的人设稳定性:'+str({'name':name,'age':age,'gender':gender,'job':job}))
    worksheet.append(output_tmp_list)
    # worksheet.append(output_time_list)
    workbook.save(OUTPUT_FILE)


if __name__ == '__main__':
    global totalTime,totalQuesNum
    print('启动接口测试程序')

    print("程序开始时间：{}".format(time.ctime()))

    preTime=time.time()
    totalTime=0#程序耗时
    totalQuesNum=0 #请求次数

    init_output(OUTPUT_FILE)
    input_list = input_init(INPUT_FILE)

    wb = load_workbook(INFO_FILE)
    sheet0 = wb[wb.sheetnames[0]]
    cols = sheet0.max_column
    rows = sheet0.max_row
    for i in tqdm(range(rows - 1)):
        name = sheet0.cell(row=i + 2, column=1).value
        info_prompt = sheet0.cell(row=i + 2, column=2).value
        info_prompt = list(json.loads(info_prompt))

        if name is not None:
            # print(info_prompt)
            # print(i, type(info_prompt))
            work(input_list, name, info_prompt)

print('程序结束时间：{}'.format(time.ctime()))
print('程序处理结果保存：{}'.format(OUTPUT_FILE))
print('程序处理总耗费时间：{:.2f}秒\n调用模型次数：{}次\n模型请求平均用时：{:.4f}秒'.format(time.time()-preTime,totalQuesNum,totalTime/totalQuesNum))